import sys, os

try:
    from cdar.models.Shallow_Water_2D.ShallowWater2DModel import ShallowWaterModelClass as model
except:
    raise ImportError('The model library cannot be found!')
        
try:
    import cdar.models.Shallow_Water_2D.ShallowWater2DModelVis as vis
except:
    raise ImportError('The vis library cannot be found!')

try:    
    from cdar.utilities import namelist_read as nml
except:
    raise ImportError('The nml library cannot be found!')

#from cdar.models.Lorenz_63.Lorenz63 import Lorenz63ModelClass as model
#import cdar.models.Lorenz_63.Lorenz63Vis as vis

try:
    import numpy as np
except:
    raise ImportError('The numpy library cannot be found!')

try:
    import ESMF
except:
    raise ImportError('The ESMP library cannot be found!')
    
SampleTime = 1800

file_path='./data'

# Read in parameters for model configurations
params = nml.namelist_read(file_path)

# start up ESMF
# this call is not necessary unless you want to to override the
# default options:
#  LogKind = NONE
#  debug = False
manager = ESMF.Manager(logkind=ESMF.LogKind.MULTI, debug=True)

# inquire for rank and proc from ESMF Virtual Machine
localPet = manager.local_pet
petCount = manager.pet_count
print "localPet = %d and petCount = %d " % (localPet, petCount)

# opening remarks
if localPet == 0:
    print "Welcome to the Gen_FG !"
           
# Set up the model
domain = model(params)

# Print domain info 
if localPet == 0:
    print domain  
    
# Perterbation FG 
domain.model_ic_add_noise()
domain.ic = domain.ic_perturbation

# Propogate the model to Sample time
output = domain.model_integration()

# Ploting and Saving
u,v,h = np.split(output[np.where(domain.history == SampleTime),:].reshape(domain.ic.shape),domain.VarDim)

if localPet == 0:
    vis.plot_3d(domain.x, domain.y, h, 'FG H at Time ='+str(SampleTime))
    np.savez(os.path.join(file_path,'FG_'+domain.Name+'_'+str(SampleTime)), FG=output[np.where(domain.history == SampleTime),:].reshape(domain.ic.shape))
    print "First guess has been saved in file: %s" %  os.path.join(file_path,'FG_'+domain.Name+'_'+str(SampleTime))+'.npz'

"""

domain = model(0,4,1001,time_diff_choice='rk4')
print domain

domain2 = model(0,4,1001,time_diff_choice='rk4')
domain2.model_ic_add_noise()
domain2.ic = domain2.ic_perturbation
[domain2.x,domain2.y,domain2.z] = np.vsplit(domain2.ic,domain2.VarDim)
print domain2

output = domain.model_integration()
output2 = domain2.model_integration()

# Pack the arraries on 3rd axis [Ntimestep, VarDim, N_trajectories]
x_t = np.dstack((output,output2))

# Rearrange the array to [N_trajectories, Ntimestep, VarDim]
x_t = np.swapaxes(np.swapaxes(x_t,2,0),1,2)

vis.movie(x_t,N_trajectories=2)
"""
